package main

import (
	"context"
	"crypto/rand"
	"database/sql"
	"encoding/json"
	"errors"
	"fmt"
	"math"
	"net/http"
	"os"
	"os/signal"
	"path/filepath"
	"sync"
	"syscall"
	"time"

	_ "github.com/go-sql-driver/mysql"
	"github.com/go-ini/ini"
	"github.com/go-redis/redis/v8"
	"github.com/gorilla/websocket"
	"github.com/sirupsen/logrus"
)

// 全局日志记录器
var logger = logrus.New()

// min 函数用于获取两个整数中的最小值
func min(a, b int) int {
	if a < b {
		return a
	}
	return b
}

// 初始化日志系统
func initLogger(level string) {
	logger.SetFormatter(&logrus.JSONFormatter{
		TimestampFormat: "2006-01-02 15:04:05",
	})

	// 根据配置设置日志级别
	logLevel, err := logrus.ParseLevel(level)
	if err != nil {
		logLevel = logrus.InfoLevel // 默认日志级别
	}
	logger.SetLevel(logLevel)
}

var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
	CheckOrigin: func(r *http.Request) bool {
		return true
	},
}

// 默认重试配置常量和WebSocket配置常量
const (
	DefaultMaxRetries    = 3
	DefaultRetryInterval = 5 * time.Second
	redisQueueKey        = "websocket_messages"
	maxMessageSize       = 1024 * 64        // 最大消息大小（字节）- 增大到64KB以支持更大的数据传输
	pongWait             = 60 * time.Second // pong消息等待时间
	writeWait            = 10 * time.Second // 写入操作超时时间
)

// 消息类型
const (
	MessageTypeData       = "data"        // 普通数据消息
	MessageTypeAck        = "ack"         // 消息确认
	MessageTypeHistoryReq = "history_req" // 历史消息请求
	MessageTypeHistoryRes = "history_res" // 历史消息响应
	MessageTypeBroadcast  = "broadcast"   // 广播消息
	MessageTypeDirect     = "direct"      // 定向消息（发送给用户所有连接）
	MessageTypeResponse   = "response"    // 响应消息
	MessageTypeError      = "error"       // 错误消息
)

// 配置结构体
type Config struct {
	ServerPassword string
	MaxConnections int
	Driver         string
	DBHost         string
	DBPort         int
	DBUsername     string
	DBPassword     string
	DBName         string
	DBCharset      string
	RedisAddr      string
	RedisPassword  string
	RedisDB        int
	ServerPort     string
	LogLevel       string
	MaxRetries     int
	RetryInterval  int // 秒
	MessageTimeout int // 消息超时时间（秒）
}

// 从配置文件加载配置
func loadConfig(configPath string) (Config, error) {
	cfg := Config{}

	// 默认配置
	cfg.ServerPassword = "myServerPassword"
	cfg.MaxConnections = 4000
	cfg.Driver = "mysql"
	cfg.DBHost = "localhost"
	cfg.DBPort = 3306
	cfg.DBUsername = "root"
	cfg.DBPassword = ""
	cfg.DBName = "websocket"
	cfg.DBCharset = "utf8mb4"
	cfg.RedisAddr = "localhost:6379"
	cfg.RedisPassword = ""
	cfg.RedisDB = 0
	cfg.ServerPort = ":8080"
	cfg.LogLevel = "info"
	cfg.MaxRetries = DefaultMaxRetries
	cfg.RetryInterval = int(DefaultRetryInterval.Seconds())
	cfg.MessageTimeout = 60 // 默认消息超时时间为60秒

	// 检查配置文件是否存在
	if _, err := os.Stat(configPath); os.IsNotExist(err) {
		logger.WithField("config_path", configPath).Warn("配置文件不存在，使用默认配置")
		return cfg, nil
	}

	// 读取配置文件
	iniCfg, err := ini.Load(configPath)
	if err != nil {
		logger.WithError(err).Warn("读取配置文件失败，使用默认配置")
		return cfg, err
	}

	// 解析配置
	if section, err := iniCfg.GetSection("Server"); err == nil {
		cfg.ServerPassword = section.Key("ServerPassword").MustString(cfg.ServerPassword)
		cfg.MaxConnections = section.Key("MaxConnections").MustInt(cfg.MaxConnections)
		cfg.ServerPort = section.Key("ServerPort").MustString(cfg.ServerPort)
		cfg.MaxRetries = section.Key("MaxRetries").MustInt(cfg.MaxRetries)
		cfg.RetryInterval = section.Key("RetryInterval").MustInt(cfg.RetryInterval)
		cfg.MessageTimeout = section.Key("MessageTimeout").MustInt(cfg.MessageTimeout)
	}

	if section, err := iniCfg.GetSection("Database"); err == nil {
		cfg.Driver = section.Key("Driver").MustString(cfg.Driver)
		cfg.DBHost = section.Key("Host").MustString(cfg.DBHost)
		cfg.DBPort = section.Key("Port").MustInt(cfg.DBPort)
		cfg.DBUsername = section.Key("Username").MustString(cfg.DBUsername)
		cfg.DBPassword = section.Key("Password").MustString(cfg.DBPassword)
		cfg.DBName = section.Key("Name").MustString(cfg.DBName)
		cfg.DBCharset = section.Key("Charset").MustString(cfg.DBCharset)
	}

	if section, err := iniCfg.GetSection("Redis"); err == nil {
		cfg.RedisAddr = section.Key("Addr").MustString(cfg.RedisAddr)
		cfg.RedisPassword = section.Key("Password").MustString(cfg.RedisPassword)
		cfg.RedisDB = section.Key("DB").MustInt(cfg.RedisDB)
	}

	if section, err := iniCfg.GetSection("Logging"); err == nil {
		cfg.LogLevel = section.Key("LogLevel").MustString(cfg.LogLevel)
	}

	logger.WithFields(logrus.Fields{
		"config_path": configPath,
		"server_port": cfg.ServerPort,
		"db_driver":   cfg.Driver,
		"db_host":     cfg.DBHost,
		"db_port":     cfg.DBPort,
		"db_name":     cfg.DBName,
	}).Info("配置文件加载成功")

	return cfg, nil
}

// Client表示一个WebSocket客户端连接
type Client struct {
	conn     *websocket.Conn // WebSocket连接
	userCode string          // 用户唯一标识
	key      string          // 用户密钥
	send     chan Message    // 待发送消息通道
	sendMu   sync.Mutex      // 保护发送操作的互斥锁，防止并发写入错误
}

// writePump方法负责将消息从send通道写入到WebSocket连接
func (c *Client) writePump(hub *Hub) {
	ticker := time.NewTicker(pongWait / 2) // 发送ping消息的定时器
	defer func() {
		ticker.Stop()
		c.conn.Close()
	}()

	for {
		select {
		case message, ok := <-c.send:
			// 检查通道是否已关闭
			if !ok {
				// 通道已关闭，发送关闭消息
				c.sendMu.Lock()
				c.conn.WriteMessage(websocket.CloseMessage, []byte{})
				c.sendMu.Unlock()
				return
			}

			// 发送消息到WebSocket连接
			c.sendMu.Lock()
			c.conn.SetWriteDeadline(time.Now().Add(writeWait))
			if err := c.conn.WriteJSON(message); err != nil {
				logger.WithError(err).WithField("userCode", c.userCode).Error("写入消息失败")
				c.sendMu.Unlock()
				// 写入失败时注销客户端
				hub.unregister <- c
				return
			}
			c.sendMu.Unlock()

		case <-ticker.C:
			// 定期发送ping消息保持连接活跃
			c.sendMu.Lock()
			c.conn.SetWriteDeadline(time.Now().Add(writeWait))
			if err := c.conn.WriteMessage(websocket.PingMessage, nil); err != nil {
				logger.WithError(err).WithField("userCode", c.userCode).Error("发送ping消息失败")
				c.sendMu.Unlock()
				// ping失败时注销客户端
				hub.unregister <- c
				return
			}
			c.sendMu.Unlock()
		}
	}
}

// readPump方法负责从WebSocket连接读取消息并处理
func (c *Client) readPump(hub *Hub) {
	defer func() {
		hub.unregister <- c
		c.conn.Close()
	}()

	for {
		var message Message
		if err := c.conn.ReadJSON(&message); err != nil {
			// 检查是否为连接关闭错误
			if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
				logger.WithError(err).WithField("userCode", c.userCode).Error("读取消息失败")
			} else {
				logger.WithField("userCode", c.userCode).Info("客户端断开连接")
			}
			return
		}

		// 填充用户代码
		message.UserCode = c.userCode

		// 根据消息类型处理
		switch message.Type {
		case MessageTypeAck:
			// 转发确认消息到Hub
			hub.broadcast <- message
		case MessageTypeHistoryReq:
			// 转发历史消息请求到Hub
			hub.broadcast <- message
		default:
			// 记录未知消息类型
			logger.WithField("message_type", message.Type).Debug("收到客户端消息")
		}
	}
}

// 增强的消息结构体，支持消息ID和确认机制
type Message struct {
	ID             string      `json:"id"`
	Type           string      `json:"type"` // "data", "ack", "history_req", "history_res", "broadcast", "error", "system"
	UserCode       string      `json:"userCode"`
	Key            string      `json:"userKey,omitempty"`
	Message        interface{} `json:"message,omitempty"`
	Timestamp      time.Time   `json:"timestamp"`
	// ServerPassword字段仅用于服务端内部验证，需要能从客户端请求中解析
	ServerPassword string      `json:"serverPassword,omitempty"` // 从JSON中的serverPassword字段映射
}

// 待确认消息结构
type PendingMessage struct {
	Message     Message
	Retries     int
	LastAttempt time.Time
	Client      *Client
}

// Hub管理WebSocket客户端连接和消息广播
type Hub struct {
	clients        map[*Client]bool
	broadcast      chan Message
	register       chan *Client
	unregister     chan *Client
	users          map[string]*Client
	mu             sync.RWMutex
	count          int
	maxConnections int
	redisCtx       context.Context
	cancelFunc     context.CancelFunc
	redisClient    *redis.Client
	isRunning      bool
	muIsRunning    sync.Mutex
	pendingMsgs    map[string]*PendingMessage
	pendingMu      sync.Mutex
	db             *sql.DB
	config         Config
    
	// 用户连接映射，用于快速查找特定用户的连接
	userConnections map[string][]*Client // 用户代码到客户端连接的映射
	userMu          sync.RWMutex         // 保护userConnections的读写锁
}

// 创建新的Hub实例
func newHub(maxConnections int, redisClient *redis.Client, db *sql.DB, config Config) *Hub {
	redisCtx, cancelFunc := context.WithCancel(context.Background())
	return &Hub{
		broadcast:       make(chan Message, 100), // 增加缓冲大小
		register:        make(chan *Client, 100), // 增加缓冲大小
		unregister:      make(chan *Client, 100), // 增加缓冲大小
		clients:         make(map[*Client]bool),
		users:           make(map[string]*Client),
		maxConnections:  maxConnections,
		redisCtx:        redisCtx,
		cancelFunc:      cancelFunc,
		redisClient:     redisClient,
		isRunning:       false,
		userConnections: make(map[string][]*Client), // 初始化用户连接映射
		pendingMsgs:     make(map[string]*PendingMessage),
		db:          db,
		config:      config,
	}
}

// 处理消息确认
func (h *Hub) handleAck(ackMsg Message) {
	// 安全的类型断言
	msgID, ok := ackMsg.Message.(string)
	if !ok {
		logger.WithField("message_data", ackMsg.Message).Error("消息确认格式错误")
		return
	}

	h.pendingMu.Lock()
	delete(h.pendingMsgs, msgID)
	h.pendingMu.Unlock()
	logger.WithFields(logrus.Fields{
		"message_id": msgID,
		"user_code":  ackMsg.UserCode,
	}).Debug("收到消息确认")

	// 更新数据库中消息的接收状态
	updateQuery := "UPDATE messages SET received = TRUE, received_at = ? WHERE id = ?"
	result, err := h.db.Exec(updateQuery, time.Now(), msgID)
	if err != nil {
		logger.WithError(err).WithField("message_id", msgID).Error("更新消息接收状态失败")
		return
	}

	// 检查更新的行数，确保至少有一行被更新
	affectedRows, err := result.RowsAffected()
	if err != nil {
		logger.WithError(err).WithField("message_id", msgID).Error("获取更新行数失败")
	} else if affectedRows == 0 {
		logger.WithField("message_id", msgID).Warn("未找到对应的消息记录进行状态更新")
	}
}

// 添加待确认消息并尝试立即发送
func (h *Hub) addPendingMessage(msg Message, client *Client) {
	// 确保消息有ID，如果没有则生成
	if msg.ID == "" {
		msg.ID = generateMessageID()
	}

	// 设置当前时间戳
	msg.Timestamp = time.Now()

	// 保存消息到数据库
	h.saveMessageToDB(msg)

	h.pendingMu.Lock()
	h.pendingMsgs[msg.ID] = &PendingMessage{
		Message:     msg,
		Retries:     0,
		LastAttempt: time.Now(),
		Client:      client,
	}
	h.pendingMu.Unlock()

	// 立即尝试发送消息，而不是等待重试计时器
	select {
	case client.send <- msg:
		logger.WithFields(logrus.Fields{
			"message_id": msg.ID,
			"userCode":   msg.UserCode,
		}).Debug("立即发送消息")
	default:
		// 通道已满，记录错误但不中断流程
		logger.WithFields(logrus.Fields{
			"message_id": msg.ID,
			"userCode":   msg.UserCode,
		}).Error("客户端消息通道已满，无法立即发送消息")
	}

	// 启动重试计时器，仅用于处理未确认的消息
	go h.startRetryTimer(msg.ID)
}

// 启动消息重试计时器
func (h *Hub) startRetryTimer(msgID string) {
	// 使用配置中的重试间隔
	ticker := time.NewTicker(time.Duration(h.config.RetryInterval) * time.Second)
	defer ticker.Stop()

	retries := 0
	maxRetries := h.config.MaxRetries

	for {
		select {
		case <-ticker.C:
			h.pendingMu.Lock()
			pendingMsg, exists := h.pendingMsgs[msgID]
			h.pendingMu.Unlock()

			if !exists {
				return // 消息已确认，退出计时器
			}

			if retries >= maxRetries {
				// 达到最大重试次数，记录错误并清理
				logger.WithField("message_id", msgID).Error("消息重试次数超过限制，放弃重试")
				h.pendingMu.Lock()
				delete(h.pendingMsgs, msgID)
				h.pendingMu.Unlock()
				return
			}

			message := pendingMsg.Message
			// 使用PendingMessage中保存的Client引用，而不是重新查找
			client := pendingMsg.Client
			// 验证客户端是否仍然存在
			h.mu.RLock()
			_, clientExists := h.clients[client]
			h.mu.RUnlock()

			if clientExists {
				// 尝试再次发送消息
				retries++
				h.pendingMu.Lock()
				pendingMsg.Retries = retries
				h.pendingMu.Unlock()

				select {
				case client.send <- message:
					logger.WithFields(logrus.Fields{
						"messageId": msgID,
						"userCode":  message.UserCode,
						"retry":     pendingMsg.Retries,
					}).Warn("重试发送消息")
				default:
					// 通道已满或已关闭，记录错误但不中断流程
					logger.WithField("messageId", msgID).Error("无法发送重试消息到客户端通道")
				}
			} else {
				logger.WithField("messageId", msgID).Warn("尝试向不存在的客户端重试发送消息")
				// 清理待确认消息，避免内存泄漏
				h.pendingMu.Lock()
				delete(h.pendingMsgs, msgID)
				h.pendingMu.Unlock()
				return
			}
		case <-time.After(time.Duration(h.config.MessageTimeout) * time.Second):
			// 消息超时，清理并退出
			logger.WithField("message_id", msgID).Error("消息超时，放弃重试")
			h.pendingMu.Lock()
			delete(h.pendingMsgs, msgID)
			h.pendingMu.Unlock()
			return
		}
	}
}

// 生成消息ID - 使用更安全的算法确保唯一性
func generateMessageID() string {
	// 使用加密随机数增强唯一性，降低碰撞概率
	b := make([]byte, 8)
	_, err := rand.Read(b)
	if err != nil {
		// 如果随机数生成失败，回退到使用时间戳和进程ID
		return fmt.Sprintf("msg-%d-%d-%d", time.Now().UnixNano(), os.Getpid(), time.Now().UnixNano()%10000)
	}
	// 结合时间戳、进程ID和随机数生成唯一ID
	return fmt.Sprintf("msg-%d-%d-%x", time.Now().UnixNano(), os.Getpid(), b)
}

// 保存消息到数据库
func (h *Hub) saveMessageToDB(msg Message) {
	messageData, err := json.Marshal(msg.Message)
	if err != nil {
		logger.WithError(err).Error("序列化消息失败")
		return
	}

	// 使用事务确保数据一致性
	tx, err := h.db.Begin()
	if err != nil {
		logger.WithError(err).Error("开始数据库事务失败")
		return
	}

	// 使用defer确保事务总是被关闭，无论执行路径如何
	defer func() {
		// 如果事务还在进行中（没有提交成功），则回滚
		if tx != nil {
			if rollbackErr := tx.Rollback(); rollbackErr != nil {
				logger.WithError(rollbackErr).Error("事务回滚失败")
			}
		}
	}()

	_, err = tx.Exec(
		"INSERT INTO messages (id, user_code, message_data, timestamp) VALUES (?, ?, ?, ?)",
		msg.ID,
		msg.UserCode,
		messageData,
		msg.Timestamp,
	)

	if err != nil {
		logger.WithError(err).WithFields(logrus.Fields{
			"message_id": msg.ID,
			"user_code":  msg.UserCode,
		}).Error("保存消息到数据库失败")
		return
	}

	if err := tx.Commit(); err != nil {
		logger.WithError(err).WithFields(logrus.Fields{
			"message_id": msg.ID,
			"user_code":  msg.UserCode,
		}).Error("提交数据库事务失败")
		return
	}

	// 提交成功后标记事务已完成
	tx = nil
}

// 获取历史消息
func (h *Hub) getHistoryMessages(userCode string, limit int, beforeID string) ([]Message, error) {
	var query string
	var args []interface{}
	var rows *sql.Rows
	var err error

	if beforeID != "" {
		query = "SELECT id, user_code, message_data, timestamp FROM messages WHERE user_code = ? AND id < ? ORDER BY timestamp DESC LIMIT ?"
		args = append(args, userCode, beforeID, limit)
	} else {
		query = "SELECT id, user_code, message_data, timestamp FROM messages WHERE user_code = ? ORDER BY timestamp DESC LIMIT ?"
		args = append(args, userCode, limit)
	}

	rows, err = h.db.Query(query, args...)
	if err != nil {
		logger.WithError(err).WithFields(logrus.Fields{
			"user_code": userCode,
		}).Error("查询历史消息失败")
		return nil, err
	}
	defer rows.Close()

	var messages []Message
	for rows.Next() {
		var msg Message
		var messageData []byte
		var timestamp time.Time

		if err := rows.Scan(&msg.ID, &msg.UserCode, &messageData, &timestamp); err != nil {
			logger.WithError(err).Error("扫描历史消息结果失败")
			continue
		}

		var data interface{}
		if err := json.Unmarshal(messageData, &data); err != nil {
			logger.WithError(err).Error("反序列化历史消息失败")
			continue
		}

		msg.Message = data
		msg.Timestamp = timestamp
		msg.Type = MessageTypeHistoryRes
		messages = append(messages, msg)
	}

	// 检查是否有行遍历错误
	if err := rows.Err(); err != nil {
		logger.WithError(err).Error("遍历历史消息结果集失败")
		return nil, err
	}

	return messages, nil
}

// Hub核心运行逻辑
func (h *Hub) run() {
	for {
		select {
		case client := <-h.register:
			h.mu.Lock()
			if h.count >= h.maxConnections {
				logger.WithField("userCode", client.userCode).Info("已达到最大连接数，关闭连接")
				errMsg := Message{Message: "已达到最大连接数"}
				client.conn.WriteJSON(errMsg)
				client.conn.Close()
				h.mu.Unlock()
				continue
			}
			h.clients[client] = true
			// 对于同一用户的多个连接，只保留第一个作为主连接
			if _, exists := h.users[client.userCode]; !exists {
				h.users[client.userCode] = client
			}
			h.count++
			logger.WithFields(logrus.Fields{
				"userCode": client.userCode,
				"totalConnections": h.count,
			}).Info("客户端已注册")
			h.mu.Unlock()

			// 更新用户连接映射（在锁外执行，避免长时间持有锁）
			h.userMu.Lock()
			if _, exists := h.userConnections[client.userCode]; !exists {
				h.userConnections[client.userCode] = make([]*Client, 0)
			}
			h.userConnections[client.userCode] = append(h.userConnections[client.userCode], client)
			h.userMu.Unlock()

		case client := <-h.unregister:
			h.mu.Lock()
			if _, ok := h.clients[client]; ok {
				delete(h.clients, client)
				// 注意：不要直接从users中删除，因为一个用户可能有多个连接
				// delete(h.users, client.userCode)
				close(client.send)
				h.count--
				logger.WithFields(logrus.Fields{
					"userCode": client.userCode,
					"totalConnections": h.count,
				}).Info("客户端断开连接")
			}
			h.mu.Unlock()

			// 从用户连接映射中移除（在锁外执行，避免长时间持有锁）
			h.userMu.Lock()
			if connections, exists := h.userConnections[client.userCode]; exists {
				// 过滤掉要移除的客户端
				newConnections := make([]*Client, 0, len(connections)-1)
				for _, c := range connections {
					if c != client {
						newConnections = append(newConnections, c)
					}
				}
				if len(newConnections) > 0 {
					h.userConnections[client.userCode] = newConnections
					// 更新users映射为剩余连接中的第一个，确保始终有有效的连接引用
					h.mu.Lock()
					h.users[client.userCode] = newConnections[0]
					h.mu.Unlock()
				} else {
					// 如果该用户没有连接了，删除用户连接映射
					delete(h.userConnections, client.userCode)
					// 同时删除users映射
					h.mu.Lock()
					delete(h.users, client.userCode)
					h.mu.Unlock()
					
					// 清理该用户以及特定客户端的所有待确认消息，避免向不存在的客户端发送消息
					h.pendingMu.Lock()
					for msgID, pendingMsg := range h.pendingMsgs {
						// 清理该用户的所有消息，或与当前断开的客户端直接关联的消息
						if pendingMsg.Message.UserCode == client.userCode || pendingMsg.Client == client {
							delete(h.pendingMsgs, msgID)
							logger.WithFields(logrus.Fields{
								"messageId": msgID,
								"userCode": client.userCode,
							}).Debug("客户端断开连接，清理关联的待确认消息")
						}
					}
					h.pendingMu.Unlock()
				}
			}
			h.userMu.Unlock()

		case message := <-h.broadcast:
			// 复制消息并清除敏感信息
			messageToSend := message
			messageToSend.ServerPassword = "" // 清除 serverPassword 字段

			// 处理不同类型的消息
			switch messageToSend.Type {
			case MessageTypeAck:
				h.handleAck(messageToSend)
			case MessageTypeHistoryReq:
				h.handleHistoryRequest(messageToSend)
			case MessageTypeDirect:
				// 处理定向消息（发送给指定用户的所有连接）
				h.sendDirectMessage(messageToSend)
			case MessageTypeData, MessageTypeResponse, MessageTypeBroadcast, "":
				// 处理标准数据消息、响应消息和广播消息
				if messageToSend.UserCode == "" {
					// 广播消息给所有客户端
					h.broadcastToAll(messageToSend)
				} else {
					// 发送消息给特定用户的特定连接（使用users映射，通常是第一个连接）
					h.sendToSpecificConnection(messageToSend)
				}
			default:
				logger.WithField("message_type", messageToSend.Type).Warn("收到未知类型的消息")
				// 默认作为数据消息处理
				if messageToSend.UserCode == "" {
					h.broadcastToAll(messageToSend)
				} else {
					h.sendToSpecificConnection(messageToSend)
				}
			}
		}
	}
}

// 处理历史消息请求
func (h *Hub) handleHistoryRequest(message Message) {
	// 安全的类型断言
	reqData, ok := message.Message.(map[string]interface{})
	if !ok {
		logger.WithField("message_data", message.Message).Error("历史消息请求格式错误")
		return
	}

	// 安全获取limit参数
	limit := 50 // 默认值
	if limitVal, exists := reqData["limit"]; exists {
		if l, ok := limitVal.(float64); ok {
			limit = int(l)
		}
	}

	// 安全获取before_id参数
	beforeID := ""
	if idVal, exists := reqData["before_id"]; exists {
		if id, ok := idVal.(string); ok {
			beforeID = id
		}
	}

	messages, err := h.getHistoryMessages(message.UserCode, limit, beforeID)
	if err != nil {
		logger.WithError(err).Error("获取历史消息失败")
		return
	}

	// 发送历史消息响应
	h.userMu.RLock()
	connections, exists := h.userConnections[message.UserCode]
	h.userMu.RUnlock()

	if exists && len(connections) > 0 {
		// 选择一个连接发送历史消息（可以是任意一个，这里选择第一个）
		for _, msg := range messages {
			connections[0].send <- msg
		}
	}
}

// 记录未找到用户的消息
func (h *Hub) recordUserNotFoundMessage(message Message) {
	// 确保消息有ID
	if message.ID == "" {
		message.ID = generateMessageID()
	}

	// 设置当前时间戳
	message.Timestamp = time.Now()

	// 记录日志
	logger.WithFields(logrus.Fields{
		"message_id": message.ID,
		"user_code": message.UserCode,
		"reason": "user_not_found",
	}).Info("记录未找到用户的消息发送尝试")

	// 保存消息到数据库，添加user_not_found标记
	messageData, err := json.Marshal(map[string]interface{}{
		"original_message": message.Message,
		"status": "user_not_found",
		"reason": "目标用户不存在或未连接",
	})
	if err != nil {
		logger.WithError(err).Error("序列化用户未找到消息失败")
		return
	}

	// 使用事务确保数据一致性
	tx, err := h.db.Begin()
	if err != nil {
		logger.WithError(err).Error("开始数据库事务失败")
		return
	}

	// 使用defer确保事务总是被关闭，无论执行路径如何
	defer func() {
		// 如果事务还在进行中（没有提交成功），则回滚
		if tx != nil {
			if rollbackErr := tx.Rollback(); rollbackErr != nil {
				logger.WithError(rollbackErr).Error("事务回滚失败")
			}
		}
	}()

	// 保存消息到数据库，使用扩展的INSERT语句包含status字段
	_, err = tx.Exec(
		"INSERT INTO messages (id, user_code, message_data, timestamp, status) VALUES (?, ?, ?, ?, ?) ON DUPLICATE KEY UPDATE message_data = VALUES(message_data), timestamp = VALUES(timestamp), status = VALUES(status)",
		message.ID,
		message.UserCode,
		messageData,
		message.Timestamp,
		"user_not_found",
	)

	if err != nil {
		// 如果使用status字段失败，尝试使用原始的INSERT语句（向后兼容）
		_, err = tx.Exec(
			"INSERT INTO messages (id, user_code, message_data, timestamp) VALUES (?, ?, ?, ?) ON DUPLICATE KEY UPDATE message_data = VALUES(message_data), timestamp = VALUES(timestamp)",
			message.ID,
			message.UserCode,
			messageData,
			message.Timestamp,
		)
		if err != nil {
			logger.WithError(err).WithFields(logrus.Fields{
				"message_id": message.ID,
				"user_code": message.UserCode,
			}).Error("保存用户未找到消息到数据库失败")
			return
		}
	}

	if err := tx.Commit(); err != nil {
		logger.WithError(err).WithFields(logrus.Fields{
			"message_id": message.ID,
			"user_code": message.UserCode,
		}).Error("提交数据库事务失败")
		return
	}

	// 提交成功后标记事务已完成
	tx = nil
}

// 发送定向消息给特定用户的所有连接
func (h *Hub) sendDirectMessage(message Message) {
	// 验证服务器密码
	if message.ServerPassword != h.config.ServerPassword {
		logger.WithField("target_user", message.UserCode).Warn("服务器密码验证失败，拒绝发送定向消息")
		return
	}

	// 获取该用户的所有连接
	h.userMu.RLock()
	connections, exists := h.userConnections[message.UserCode]
	// 创建连接副本，避免在锁外引用
	connectionsCopy := make([]*Client, len(connections))
	copy(connectionsCopy, connections)
	h.userMu.RUnlock()

	if !exists || len(connectionsCopy) == 0 {
		logger.WithField("user_code", message.UserCode).Warn("未找到用户的任何连接，记录消息发送尝试")
		// 记录未找到用户的消息
		h.recordUserNotFoundMessage(message)
		return
	}

	logger.WithFields(logrus.Fields{
		"userCode": message.UserCode,
		"connCount": len(connectionsCopy),
	}).Debug("发送定向消息给用户的所有连接")

	// 向每个连接发送消息
	for _, client := range connectionsCopy {
		go h.addPendingMessage(message, client)
	}
}

// 广播消息给所有客户端
func (h *Hub) broadcastToAll(message Message) {
	h.mu.RLock()
	// 创建客户端副本，避免在锁外引用
	clientsCopy := make([]*Client, 0, len(h.clients))
	for client := range h.clients {
		clientsCopy = append(clientsCopy, client)
	}
	h.mu.RUnlock()

	logger.WithField("client_count", len(clientsCopy)).Debug("广播消息给所有客户端")

	// 向每个客户端发送消息，为每个客户端创建独立的消息ID
	for _, client := range clientsCopy {
		// 为每个客户端创建一个消息副本
		clientMsg := message
		// 不使用原始消息的ID，而是在addPendingMessage中为每个客户端生成新的独立ID
		clientMsg.ID = ""
		go h.addPendingMessage(clientMsg, client)
	}
}

// 发送消息给特定用户的特定连接（通常是第一个连接）
func (h *Hub) sendToSpecificConnection(message Message) {
	h.mu.RLock()
	client, ok := h.users[message.UserCode]
	// 验证连接是否存在
	h.mu.RUnlock()

	if !ok {
		logger.WithField("user_code", message.UserCode).Warn("未找到用户，记录消息发送尝试")
		// 记录未找到用户的消息
		h.recordUserNotFoundMessage(message)
		return
	}

	logger.WithField("user_code", message.UserCode).Debug("发送消息给特定客户端")

	// 发送消息
	go h.addPendingMessage(message, client)
}

// 重新连接Redis
func reconnectRedis(hub *Hub) error {
	logger.Info("尝试重新连接Redis...")
	// 创建一个新的Redis客户端
	newClient := redis.NewClient(&redis.Options{
		Addr:     hub.config.RedisAddr,
		Password: hub.config.RedisPassword,
		DB:       hub.config.RedisDB,
	})

	// 测试新连接
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	if _, err := newClient.Ping(ctx).Result(); err != nil {
		return err
	}

	// 原子地替换旧的客户端 - 使用Lock而不是RLock，因为这是写操作
	hub.mu.Lock()
	hub.redisClient = newClient
	hub.mu.Unlock()

	return nil
}

// 判断错误是否为Redis连接错误
type redisConnError interface {
	Timeout() bool
	Temporary() bool
}

func isRedisConnectionError(err error) bool {
	var connErr redisConnError
	ok := errors.As(err, &connErr)
	if !ok {
		return false
	}
	// 连接错误通常是临时性的，但不是超时错误
	return connErr.Temporary() && !connErr.Timeout()
}

// WebSocket连接处理函数
func serveWs(hub *Hub, w http.ResponseWriter, r *http.Request) {
	// 验证请求速率
	clientIP, isValid := validateRequest(r)
	if !isValid {
		http.Error(w, "请求过于频繁，请稍后再试", http.StatusTooManyRequests)
		return
	}

	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		logger.WithError(err).WithField("ip", clientIP).Error("WebSocket升级失败")
		return
	}

	// 设置连接的写入和读取超时
	conn.SetReadLimit(maxMessageSize)
	conn.SetReadDeadline(time.Now().Add(pongWait))
	conn.SetPongHandler(func(string) error {
		conn.SetReadDeadline(time.Now().Add(pongWait))
		return nil
	})

	// 读取初始消息（包含认证信息）
	var msg Message
	if err := conn.ReadJSON(&msg); err != nil {
		logger.WithError(err).WithField("ip", clientIP).Error("读取初始认证消息失败")
		// 尝试发送错误消息并正确关闭连接
		conn.WriteControl(websocket.CloseMessage,
			websocket.FormatCloseMessage(websocket.CloseProtocolError, "读取认证信息失败"),
			time.Now().Add(writeWait))
		conn.Close()
		return
	}

	// 验证WebSocket认证信息
	if !validateWebSocketAuth(msg) {
		errMsg := Message{Type: "error", Message: "WebSocket认证失败"}
		conn.SetWriteDeadline(time.Now().Add(writeWait))
		conn.WriteJSON(errMsg)
		conn.Close()
		return
	}

	// 验证服务器密码
	if msg.ServerPassword != hub.config.ServerPassword {
		logger.WithField("userCode", msg.UserCode).Warn("服务器密码无效")
		errMsg := Message{Type: "error", Message: "服务器密码无效"}
		// 设置写入超时
		conn.SetWriteDeadline(time.Now().Add(writeWait))
		if err := conn.WriteJSON(errMsg); err != nil {
			logger.WithError(err).WithField("userCode", msg.UserCode).Error("发送密码验证错误消息失败")
		}
		conn.Close()
		return
	}

	// 验证用户代码和密钥不为空
	if msg.UserCode == "" {
		logger.Warn("用户代码为空")
		errMsg := Message{Type: "error", Message: "用户代码不能为空"}
		conn.SetWriteDeadline(time.Now().Add(writeWait))
		conn.WriteJSON(errMsg)
		conn.Close()
		return
	}

	// 创建客户端并注册
	client := &Client{conn: conn, userCode: msg.UserCode, key: msg.Key, send: make(chan Message, 256)}
	hub.register <- client

	// 启动读写泵协程
	go client.writePump(hub)
	go client.readPump(hub)

	// 发送登录成功消息
	successMsg := Message{Type: "system", Message: "登录成功"}
	conn.SetWriteDeadline(time.Now().Add(writeWait))
	if err := conn.WriteJSON(successMsg); err != nil {
		logger.WithError(err).WithField("user_code", msg.UserCode).Error("发送登录成功消息失败")
		// 失败时也继续，因为连接已经注册
	}
}

// HTTP API处理函数 - 发送消息
func sendMessageHandler(hub *Hub) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		// 验证请求速率
		clientIP, isValid := validateRequest(r)
		if !isValid {
			http.Error(w, "请求过于频繁，请稍后再试", http.StatusTooManyRequests)
			logger.WithField("ip", clientIP).Warn("API请求超过速率限制")
			return
		}

		w.Header().Set("Content-Type", "application/json")
		var msg Message
		if err := json.NewDecoder(r.Body).Decode(&msg); err != nil {
			logger.WithError(err).WithField("ip", clientIP).Error("解析消息体失败")
			w.WriteHeader(http.StatusBadRequest)
			json.NewEncoder(w).Encode(map[string]string{"error": "请求体格式错误"})
			return
		}

		// 验证服务器密码
		if msg.ServerPassword != hub.config.ServerPassword {
			logger.WithField("userCode", msg.UserCode).WithField("ip", clientIP).Warn("服务器密码无效")
			w.WriteHeader(http.StatusUnauthorized)
			json.NewEncoder(w).Encode(map[string]string{"error": "服务器密码无效"})
			return
		}

		// 验证特定用户的消息
		if msg.UserCode != "" {
			hub.mu.RLock()
			c, ok := hub.users[msg.UserCode]
			validKey := ok && c.key == msg.Key
			hub.mu.RUnlock()

			if !ok {
				logger.WithField("userCode", msg.UserCode).WithField("ip", clientIP).Warn("未找到用户，记录消息发送尝试")
				// 记录未找到用户的消息
				hub.recordUserNotFoundMessage(msg)
				w.WriteHeader(http.StatusNotFound)
				json.NewEncoder(w).Encode(map[string]string{"error": "未找到用户，记录消息发送尝试"})
				return
			}
			if !validKey {
				logger.WithField("userCode", msg.UserCode).WithField("ip", clientIP).Warn("用户密钥无效")
				w.WriteHeader(http.StatusUnauthorized)
				json.NewEncoder(w).Encode(map[string]string{"error": "用户密钥无效"})
				return
			}
		}

		// 确保消息有ID，如果没有则生成
		if msg.ID == "" {
			msg.ID = generateMessageID()
		}

		// 发送消息到广播通道
		hub.broadcast <- msg
		logger.WithField("userCode", msg.UserCode).WithField("ip", clientIP).Info("消息发送成功")
		w.WriteHeader(http.StatusOK)
		json.NewEncoder(w).Encode(map[string]string{"status": "success", "message": "发送成功", "messageId": msg.ID})
	}
}

// HTTP API处理函数 - 获取用户连接数
func getUserCountHandler(hub *Hub) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		hub.mu.RLock()
		defer hub.mu.RUnlock()

		count := len(hub.clients)
		userCodes := make([]string, 0, count)
		for userCode := range hub.users {
			userCodes = append(userCodes, userCode)
		}

		response := struct {
			Count     int      `json:"count"`
			UserCodes []string `json:"user_codes"`
		}{}
		response.Count = count
		response.UserCodes = userCodes

		json.NewEncoder(w).Encode(response)
	}
}

// HTTP API处理函数 - 获取用户连接详情（包含每个用户的连接数）
func getUserConnectionsHandler(hub *Hub) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")

		// 验证服务器密码
		serverPassword := r.URL.Query().Get("server_password")
		if serverPassword != hub.config.ServerPassword {
			logger.Warn("获取用户连接详情：服务器密码验证失败")
			w.WriteHeader(http.StatusUnauthorized)
			json.NewEncoder(w).Encode(map[string]string{"error": "服务器密码无效"})
			return
		}

		hub.userMu.RLock()
		defer hub.userMu.RUnlock()

		// 构建响应数据
		connectionsInfo := make(map[string]int)
		totalConnections := 0

		for userCode, connections := range hub.userConnections {
			connectionsInfo[userCode] = len(connections)
			totalConnections += len(connections)
		}

		response := struct {
			TotalConnections int              `json:"total_connections"`
			UserConnections  map[string]int   `json:"user_connections"`
			Users            []string         `json:"users"`
		}{}
		response.TotalConnections = totalConnections
		response.UserConnections = connectionsInfo

		// 提取所有用户名
		users := make([]string, 0, len(connectionsInfo))
		for userCode := range connectionsInfo {
			users = append(users, userCode)
		}
		response.Users = users

		json.NewEncoder(w).Encode(response)
	}
}

// Redis消息队列消费者
func startRedisConsumer(hub *Hub) {
	hub.muIsRunning.Lock()
	hub.isRunning = true
	hub.muIsRunning.Unlock()

	logger.Info("Redis消费者已启动")

	// 错误重试计数器和指数退避参数
	retryCount := 0
	maxRetries := 5
	baseRetryTime := 1 * time.Second

	// Redis连接检查定时器
	connectionCheckTicker := time.NewTicker(30 * time.Second)
	defer connectionCheckTicker.Stop()

	for {
		// 检查上下文是否已取消
		if hub.redisCtx.Err() != nil {
			cleanupRedisConsumer(hub, "上下文已取消，Redis消费者停止")
			return
		}

		// 检查Redis连接状态
		select {
		case <-connectionCheckTicker.C:
			if err := checkRedisConnection(hub); err != nil {
				logger.WithError(err).Warn("Redis连接状态检查失败")
				// 尝试重新连接
				if err := reconnectRedis(hub); err != nil {
					logger.WithError(err).Error("Redis重新连接失败")
				} else {
					logger.Info("Redis重新连接成功")
				}
			}

		default:
			// 继续主逻辑
		}

		result, err := hub.redisClient.BLPop(hub.redisCtx, 0, redisQueueKey).Result()
		if err != nil {
			// 区分上下文取消错误和其他错误
			if errors.Is(err, hub.redisCtx.Err()) {
				cleanupRedisConsumer(hub, "上下文已取消，Redis消费者停止")
				return
			}

			// 专门处理连接错误
			if isRedisConnectionError(err) {
				logger.WithError(err).Error("Redis连接错误，尝试重新连接")
				// 尝试重新连接
				if err := reconnectRedis(hub); err != nil {
					logger.WithError(err).Error("Redis重新连接失败")
				} else {
					logger.Info("Redis重新连接成功")
				}
			}

			// 实现指数退避重试
			retryCount++
			logger.WithError(err).WithField("retry_count", retryCount).Error("读取Redis队列时出错")

			// 计算退避时间，最大不超过30秒
			retryTime := time.Duration(math.Min(float64(baseRetryTime)*math.Pow(2, float64(retryCount-1)), 30*float64(time.Second)))
			timer := time.NewTimer(retryTime)

			select {
			case <-hub.redisCtx.Done():
				timer.Stop()
				cleanupRedisConsumer(hub, "上下文已取消，Redis消费者停止")
				return
			case <-timer.C:
				// 重试完成，继续循环
			}

			// 如果达到最大重试次数，重置计数器
			if retryCount >= maxRetries {
				retryCount = 0
			}

			continue
		}

		// 重置重试计数器，因为操作成功
		retryCount = 0

		if len(result) > 1 {
			var msg Message
			if err := json.Unmarshal([]byte(result[1]), &msg); err != nil {
				logger.WithError(err).Error("解组Redis消息时出错")
				// 添加错误计数和监控指标
				continue
			}

			// 验证消息有效性
			if msg.UserCode == "" && msg.Type != MessageTypeBroadcast {
				logger.WithField("message_id", msg.ID).Warn("接收到无效的Redis消息：用户代码为空且非广播消息")
				continue
			}

			// 安全地将消息发送到广播通道
			sendRedisMessageToHub(hub, msg)
		}
	}
}

// 清理Redis消费者资源
func cleanupRedisConsumer(hub *Hub, reason string) {
	logger.Info(reason)
	hub.muIsRunning.Lock()
	hub.isRunning = false
	hub.muIsRunning.Unlock()
}

// 安全地将Redis消息发送到Hub广播通道
func sendRedisMessageToHub(hub *Hub, msg Message) {
	// 设置超时上下文，防止长时间阻塞
	sendCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	select {
	case hub.broadcast <- msg:
		logger.WithFields(logrus.Fields{
			"userCode": msg.UserCode,
			"type": msg.Type,
		}).Info("已向用户发送Redis消息")
	case <-hub.redisCtx.Done():
		logger.Info("上下文已取消，停止发送Redis消息")
	case <-sendCtx.Done():
			logger.WithField("userCode", msg.UserCode).Error("发送Redis消息超时")
	}
}

// HTTP API处理函数 - 控制Redis消费者
func controlRedisConsumerHandler(hub *Hub) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		action := r.URL.Query().Get("action")
		hub.muIsRunning.Lock()
		isRunning := hub.isRunning
		hub.muIsRunning.Unlock()

		switch action {
		case "start":
			if isRunning {
				logger.Warn("Redis消费者已在运行")
				http.Error(w, "Redis消费者已在运行", http.StatusBadRequest)
				return
			}
			hub.muIsRunning.Lock()
			// 双重检查，确保在锁定期间状态没有变化
			if !hub.isRunning {
				if hub.cancelFunc != nil {
					hub.cancelFunc()
					hub.cancelFunc = nil
				}
				hub.redisCtx, hub.cancelFunc = context.WithCancel(context.Background())
				go startRedisConsumer(hub)
			}
			hub.muIsRunning.Unlock()
			w.WriteHeader(http.StatusOK)
			json.NewEncoder(w).Encode("Redis消费者已启动")
		case "stop":
			if !isRunning {
				logger.Warn("Redis消费者已停止")
				http.Error(w, "Redis消费者已停止", http.StatusBadRequest)
				return
			}
			hub.muIsRunning.Lock()
			// 双重检查，确保在锁定期间状态没有变化
			if hub.isRunning && hub.cancelFunc != nil {
				hub.cancelFunc()
				hub.cancelFunc = nil
				// isRunning将由Redis消费者协程更新
			}
			hub.muIsRunning.Unlock()
			w.WriteHeader(http.StatusOK)
			json.NewEncoder(w).Encode("Redis消费者已停止")
		default:
			logger.Warnf("无效的操作: %s", action)
			http.Error(w, "无效的操作", http.StatusBadRequest)
		}
	}
}

// HTTP API处理函数 - 获取Redis消费者状态
func getRedisConsumerStatusHandler(hub *Hub) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		hub.muIsRunning.Lock()
		status := hub.isRunning
		hub.muIsRunning.Unlock()

		response := struct {
			Status bool `json:"status"`
		}{}
		response.Status = status

		json.NewEncoder(w).Encode(response)
	}
}

// 主函数
func main() {
	// 获取应用程序所在目录
	appDir, err := os.Executable()
	if err != nil {
		logger.WithError(err).Warn("无法获取应用程序目录，使用当前工作目录")
		appDir = "."
	} else {
		appDir = filepath.Dir(appDir)
	}

	// 构建配置文件路径
	configPath := filepath.Join(appDir, "config.ini")

	// 加载配置
	config, err := loadConfig(configPath)
	if err != nil {
		logger.WithError(err).Warn("配置文件加载不完全，部分配置使用默认值")
	}

	// 初始化日志系统
	initLogger(config.LogLevel)

	// 连接MySQL数据库 - 添加时区设置以确保时间戳正确处理
	connStr := fmt.Sprintf("%s:%s@tcp(%s:%d)/%s?charset=%s&parseTime=True&loc=Local",
		config.DBUsername, config.DBPassword, config.DBHost, config.DBPort, config.DBName, config.DBCharset)
	db, err := sql.Open(config.Driver, connStr)
	if err != nil {
		logger.WithError(err).Fatal("连接数据库失败")
	}
	defer db.Close()

	// 配置MySQL连接池参数，根据WebSocket服务器的高并发特性调整
	db.SetMaxOpenConns(50)                  // 最大打开连接数 - 提高以适应高并发
	db.SetMaxIdleConns(25)                  // 最大空闲连接数 - 设为最大打开连接数的一半
	db.SetConnMaxLifetime(time.Hour)        // 连接最大生命周期
	db.SetConnMaxIdleTime(time.Minute * 30) // 连接最大空闲时间

	// 测试数据库连接
	if err = db.Ping(); err != nil {
		logger.WithError(err).Fatal("测试数据库连接失败")
	}
	logger.Info("数据库连接成功")

	// 创建消息表（如果不存在）
	_, err = db.Exec(`
		CREATE TABLE IF NOT EXISTS messages (
			id VARCHAR(100) PRIMARY KEY,
			user_code VARCHAR(100) NOT NULL,
			message_data TEXT NOT NULL,
			timestamp DATETIME NOT NULL,
			received BOOLEAN DEFAULT FALSE,
			received_at DATETIME NULL,
			INDEX idx_user_timestamp (user_code, timestamp)
		)
	`)
	if err != nil {
		logger.WithError(err).Fatal("创建消息表失败")
	}
	logger.Info("消息表创建成功")

	// 连接Redis
	redisClient := redis.NewClient(&redis.Options{
		Addr:     config.RedisAddr,
		Password: config.RedisPassword,
		DB:       config.RedisDB,
	})

	// 测试Redis连接
	ctx := context.Background()
	if _, err = redisClient.Ping(ctx).Result(); err != nil {
		logger.WithError(err).Fatal("连接Redis失败")
	}
	logger.Info("Redis连接成功")

	// 创建Hub
	hub := newHub(config.MaxConnections, redisClient, db, config)

	// 确保在程序退出时清理资源
	defer func() {
		if hub.cancelFunc != nil {
			hub.cancelFunc()
		}
		// 等待Hub清理完成（实际项目中可能需要更复杂的优雅关闭机制）
		time.Sleep(500 * time.Millisecond)
	}()

	// 启动Hub
	go hub.run()

	// 注册HTTP路由
	http.HandleFunc("/ws", func(w http.ResponseWriter, r *http.Request) {
		serveWs(hub, w, r)
	})
	http.HandleFunc("/send-message", sendMessageHandler(hub))
	http.HandleFunc("/user-count", getUserCountHandler(hub))
	http.HandleFunc("/user-connections", getUserConnectionsHandler(hub)) // 新增路由：获取用户连接详情
	http.HandleFunc("/control-consumer", controlRedisConsumerHandler(hub))
	http.HandleFunc("/consumer-status", getRedisConsumerStatusHandler(hub))

	// 创建服务器实例，以便能够优雅关闭
	server := &http.Server{
		Addr:    config.ServerPort,
		Handler: nil,
	}

	// 启动HTTP服务器在单独的goroutine中
	go func() {
		logger.Info("WebSocket服务器已启动在 " + config.ServerPort)
		if err := server.ListenAndServe(); err != nil {
			// 忽略http.ErrServerClosed错误，这是预期的关闭
			if err != http.ErrServerClosed {
				logger.WithError(err).Fatal("服务器启动失败")
			}
		}
	}()

	// 设置信号处理，捕获中断信号以便优雅关闭
	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)

	// 阻塞等待信号
	<-sigChan

	logger.Info("收到关闭信号，正在优雅关闭服务器...")

	// 创建5秒超时的上下文用于服务器关闭
	shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	// 关闭HTTP服务器
	if err := server.Shutdown(shutdownCtx); err != nil {
		logger.WithError(err).Error("服务器关闭失败")
	}

	// 清理Hub资源
	if hub.cancelFunc != nil {
		hub.cancelFunc()
	}

	logger.Info("服务器已成功关闭")
}

// 验证请求函数
func validateRequest(r *http.Request) (string, bool) {
	// 获取客户端IP
	clientIP := r.RemoteAddr

	// 移除速率限制检查，始终返回true
	return clientIP, true
}

// 验证WebSocket认证消息
func validateWebSocketAuth(msg Message) bool {
	// 验证必要字段
	if msg.UserCode == "" {
		logger.Warn("WebSocket认证失败：用户代码为空")
		return false
	}

	// 验证消息结构完整性
	if len(msg.UserCode) > 100 {
		logger.WithField("user_code", msg.UserCode).Warn("WebSocket认证失败：用户代码过长")
		return false
	}

	return true
}

// 检查Redis连接状态
func checkRedisConnection(hub *Hub) error {
	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancel()
	_, err := hub.redisClient.Ping(ctx).Result()
	return err
}
